import torch
from u2net import U2NETP
from torchvision import transforms
import numpy as np
import os
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import cfg
from torchvision.utils import save_image


class Detect():
    def __init__(self, netSavePath):
        # 设定运行网络
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.net = U2NETP().to(self.device)
        # 加载参数
        if os.path.exists(netSavePath):
            self.net.load_state_dict(torch.load(netSavePath, map_location=self.device))
            print('u2net神经网络参数加载成功！')
        else:
            print('No params')

    def Generation_mask(self, src):
        w, h = src.shape
        img = np.array((Image.fromarray(src)).resize((cfg.img_size, cfg.img_size), Image.ANTIALIAS))
        # img = np.array((Image.fromarray(src)).resize((128, 128), Image.ANTIALIAS))
        img_data = transforms.ToTensor()(img)

        # 增加一维
        img_data = img_data.unsqueeze(0).to(self.device)

        # 将程序放在网络中
        d0, d1, d2, d3, d4, d5, d6 = self.net(img_data)

        # # 制作mask
        d0 = d0.cpu().detach()
        d0 = torch.where(d0 > 0.5, 1, 0)


        out_Img = (torch.cat([d0[0], d0[0] * 255, d0[0] * 255], 0)).numpy()

        mask = out_Img.transpose(1, 2, 0).astype(np.uint8)
        cv2.imshow("d0", mask)
        cv2.waitKey(0)

        print(mask.shape)
        mask = np.array((Image.fromarray(mask)).resize((h, w), Image.ANTIALIAS))

        return mask[..., 0]


if __name__ == '__main__':
    pth_path = 'params.pth'  # 网络参数

    path = r'E:\projetct\head_CT\data\20210827_labled_AAA_CTA\train_test\test.npy'
    save_path = r'E:\projetct\head_CT\data\20210827_labled_AAA_CTA\predict_imgs'  # 保存地址

    data = np.load(path)
    detect = Detect(pth_path)

    for i in range(len(data)):
        new_arr = data[i].copy()
        mask = detect.Generation_mask(new_arr[0])
        print(np.max(mask))
        cv2.imshow("mask",mask)
        cv2.waitKey(0)
        new_arr_predict=new_arr[0].copy()
        new_arr_target=new_arr[0].copy()
        # exit()

        # #################################  画图 检查掩码和图片的对应关系#######################################
        new_arr_target[new_arr[1] == 1] = np.array([2000])  # 将分割结果中是腹主动脉的位置蒙上一层浅红色
        new_arr_predict[mask == 255] = np.array([2000])  # 将分割结果中是腹主动脉的位置蒙上一层浅红色

        new_image = np.hstack((new_arr_target, new_arr[1], new_arr_predict, data[i][0]))
        # plt.imsave("{}.jpg".format(i),new_arr_predict)
        plt.figure(figsize=(10, 10))
        plt.imshow(new_image)
        plt.show()
        plt.pause(0)
